import pandas as pd
import os
from datetime import timedelta
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import numpy as np


# --- 数据加载函数 (无需修改) ---
def load_data_in_range(start_time_str: str, end_time_str: str, base_path: str, symbol: str,
                       timeframe: str) -> pd.DataFrame:
    """从指定的CSV文件中加载时间范围内的数据。"""
    start_date = pd.to_datetime(start_time_str)
    end_date = pd.to_datetime(end_time_str)
    data_folder_path = os.path.join(base_path, symbol, timeframe)
    if not os.path.isdir(data_folder_path):
        print(f"错误: 文件夹不存在 -> {data_folder_path}")
        return pd.DataFrame()

    date_range = pd.date_range(start=start_date.date(), end=end_date.date(), freq='D').strftime('%Y-%m').unique()
    if len(date_range) == 0: date_range = [start_date.strftime('%Y-%m')]
    all_dfs = []
    print(f"将在以下月份文件中查找数据: {list(date_range)}")
    for year_month in date_range:
        file_path = os.path.join(data_folder_path, f"{symbol.upper()}-{timeframe}-{year_month}.csv")
        if os.path.exists(file_path):
            print(f"正在读取文件: {file_path}")
            try:
                df = pd.read_csv(file_path)
                required_cols = ['open_time', 'open', 'high', 'low', 'close', 'volume', 'taker_buy_volume']
                if not all(col in df.columns for col in required_cols):
                    print(f"错误: 文件 {file_path} 缺少必要的列。")
                    continue
                all_dfs.append(df)
            except Exception as e:
                print(f"读取或处理文件 {file_path} 时出错: {e}")
        else:
            print(f"警告: 未找到文件 {file_path}")

    if not all_dfs:
        print("错误: 在指定时间范围内未找到任何数据文件。")
        return pd.DataFrame()

    combined_df = pd.concat(all_dfs, ignore_index=True)
    combined_df.drop_duplicates(subset=['open_time'], inplace=True)
    return combined_df


# --- 数据处理函数 (无需修改) ---
def process_data(df: pd.DataFrame, zscore_period: int) -> pd.DataFrame:
    """
    处理原始数据，包括：
    1. 转换时间到UTC+8。
    2. 计算价格变动率和成交量的Z-Score。
    """
    if df.empty:
        return df

    df['open_time'] = pd.to_datetime(df['open_time'], unit='ms', utc=True)
    df['open_time'] = df['open_time'].dt.tz_convert('Asia/Shanghai')
    df.set_index('open_time', inplace=True)

    numeric_cols = ['open', 'high', 'low', 'close', 'volume', 'taker_buy_volume']
    for col in numeric_cols:
        df[col] = pd.to_numeric(df[col], errors='coerce')
    df.dropna(subset=numeric_cols, inplace=True)
    df.sort_index(inplace=True)

    df['volume_delta'] = (2 * df['taker_buy_volume']) - df['volume']
    df['cvd'] = df['volume_delta'].cumsum()

    df['price_change_rate'] = (df['close'] - df['open']).abs() / df['close'].replace(0, np.nan)
    price_change_mean = df['price_change_rate'].rolling(window=zscore_period, min_periods=1).mean()
    price_change_std = df['price_change_rate'].rolling(window=zscore_period, min_periods=1).std()
    volume_mean = df['volume'].rolling(window=zscore_period, min_periods=1).mean()
    volume_std = df['volume'].rolling(window=zscore_period, min_periods=1).std()

    df['price_change_zscore'] = (df['price_change_rate'] - price_change_mean) / price_change_std.replace(0, np.nan)
    df['volume_zscore'] = (df['volume'] - volume_mean) / volume_std.replace(0, np.nan)
    df[['price_change_zscore', 'volume_zscore']] = df[['price_change_zscore', 'volume_zscore']].fillna(0)

    print(f"数据处理和Z-Score计算完成。共处理 {len(df)} 条K线。")
    return df


# --- 绘图函数 (已移除枢轴点) ---
def plot_simplified_zscore_chart(df: pd.DataFrame, chart_title: str, zscore_threshold: float = 2.0):
    """
    绘制简化的图表，并根据K线方向和Z-Score阈值进行多空标注。
    """
    if df.empty:
        print("数据为空，无法绘制图表。")
        return

    hover_texts = []
    for index, row in df.iterrows():
        text = (f"<b>Time (UTC+8)</b>: {index.strftime('%Y-%m-%d %H:%M')}<br>"
                f"<b>Open</b>: {row['open']:.5f}<br>"
                f"<b>High</b>: {row['high']:.5f}<br>"
                f"<b>Low</b>: {row['low']:.5f}<br>"
                f"<b>Close</b>: {row['close']:.5f}<br>"
                f"<b>Volume</b>: {row['volume']:,.0f}<br>"
                f"<b>Price Change Z-Score</b>: {row['price_change_zscore']:.5f}<br>"
                f"<b>Volume Z-Score</b>: {row['volume_zscore']:.5f}")
        hover_texts.append(text)

    fig = make_subplots(
        rows=3, cols=1,
        shared_xaxes=True,
        vertical_spacing=0.03,
        row_heights=[0.7, 0.15, 0.15]
    )

    fig.add_trace(go.Candlestick(x=df.index,
                                 open=df['open'], high=df['high'], low=df['low'], close=df['close'],
                                 name='OHLC', hoverinfo='text', hovertext=hover_texts),
                  row=1, col=1)

    # --- 根据多空方向分别标注信号 ---
    highlight_df = df[(df['price_change_zscore'] > zscore_threshold) & (df['volume_zscore'] > zscore_threshold)]

    if not highlight_df.empty:
        # 筛选看涨信号 (阳线)
        long_signals_df = highlight_df[highlight_df['close'] > highlight_df['open']]
        if not long_signals_df.empty:
            fig.add_trace(go.Scatter(
                x=long_signals_df.index,
                y=long_signals_df['low'] * 0.995,
                mode='markers+text',
                marker=dict(symbol='arrow-up', color='green', size=12),
                text=['Long'] * len(long_signals_df),
                textposition='bottom center',
                textfont=dict(color='green', size=10),
                name='Long Signal',
                hoverinfo='none'
            ), row=1, col=1)

        # 筛选看跌信号 (阴线或十字星)
        short_signals_df = highlight_df[highlight_df['close'] <= highlight_df['open']]
        if not short_signals_df.empty:
            fig.add_trace(go.Scatter(
                x=short_signals_df.index,
                y=short_signals_df['high'] * 1.005,
                mode='markers+text',
                marker=dict(symbol='arrow-down', color='red', size=12),
                text=['Short'] * len(short_signals_df),
                textposition='top center',
                textfont=dict(color='red', size=10),
                name='Short Signal',
                hoverinfo='none'
            ), row=1, col=1)

        print(f"\n在图上标注了 {len(long_signals_df)} 个看涨信号和 {len(short_signals_df)} 个看跌信号。")

    pcz_colors = ['blue' if z >= 0 else 'orange' for z in df['price_change_zscore']]
    fig.add_trace(go.Bar(x=df.index, y=df['price_change_zscore'], name='Price Change Z-Score', marker_color=pcz_colors), row=2, col=1)
    fig.add_hline(y=zscore_threshold, line_dash="dot", line_color="red", row=2, col=1)

    vz_colors = ['blue' if z >= 0 else 'orange' for z in df['volume_zscore']]
    fig.add_trace(go.Bar(x=df.index, y=df['volume_zscore'], name='Volume Z-Score', marker_color=vz_colors), row=3, col=1)
    fig.add_hline(y=zscore_threshold, line_dash="dot", line_color="red", row=3, col=1)

    start_str = df.index[0].strftime('%Y-%m-%d %H:%M')
    end_str = df.index[-1].strftime('%Y-%m-%d %H:%M')
    full_title = f'{chart_title} from {start_str} to {end_str} (UTC+8)'

    fig.update_layout(
        title_text=full_title,
        xaxis_rangeslider_visible=False,
        showlegend=False,
        yaxis1_title='Price',
        yaxis2_title='Price ZS',
        yaxis3_title='Vol ZS',
        hovermode='x unified'
    )
    fig.show()


if __name__ == '__main__':
    # ==================== 用户配置区 ====================
    BASE_DATA_PATH = "F:/personal/binance_klines"
    SYMBOL = "SUIUSDT"
    TIMEFRAME = "1h"

    CHART_START_TIME = "2024-08-01 00:00:00"
    CHART_END_TIME = "2024-09-01 00:00:00"

    PREWARM_DAYS = 10
    ZSCORE_PERIOD = 100
    ZSCORE_THRESHOLD = 2.0
    # ====================================================

    print("--- K线图表程序启动 ---")

    chart_start_dt_obj = pd.to_datetime(CHART_START_TIME)
    chart_end_dt_obj = pd.to_datetime(CHART_END_TIME)

    # --- 移除枢轴点计算 ---
    # pivot_points = {} # 不再需要

    # --- 数据加载与预热 ---
    data_load_start_dt = chart_start_dt_obj - timedelta(days=PREWARM_DAYS)
    print(f"\n--- 正在加载图表和预热数据 ({data_load_start_dt.strftime('%Y-%m-%d %H:%M')} to {chart_end_dt_obj.strftime('%Y-%m-%d %H:%M')}) ---")
    raw_data_df = load_data_in_range(data_load_start_dt.strftime('%Y-%m-%d %H:%M:%S'),
                                     chart_end_dt_obj.strftime('%Y-%m-%d %H:%M:%S'),
                                     BASE_DATA_PATH, SYMBOL, TIMEFRAME)

    # --- 数据处理 ---
    processed_df = process_data(raw_data_df, zscore_period=ZSCORE_PERIOD)

    # 筛选最终图表所需的时间范围
    chart_start_dt_aware = chart_start_dt_obj.tz_localize('Asia/Shanghai')
    chart_end_dt_aware = chart_end_dt_obj.tz_localize('Asia/Shanghai')
    final_df = processed_df.loc[chart_start_dt_aware:chart_end_dt_aware]

    print(f"\n数据预热和处理完成，最终图表将显示 {len(final_df)} 条K线。")

    # --- 绘图 ---
    if not final_df.empty:
        chart_title_str = f"{SYMBOL} - {TIMEFRAME}"
        plot_simplified_zscore_chart(final_df,
                                     chart_title=chart_title_str,
                                     zscore_threshold=ZSCORE_THRESHOLD)
        print("\n--- 图表生成完毕 ---")
    else:
        print("\n未能加载或处理任何图表数据，程序退出。请检查您的配置和文件路径。")